
import os
import sys

sys.path.append(os.path.split(sys.path[0])[0])
import shutil
from time import time
import SimpleITK as sitk
import scipy.ndimage as ndimage
from multiprocessing.dummy import Pool
import math
import numpy as np
'''
裁剪MRI的中心一点的体积，比如100*512*512，才建成100*256*256
'''

data_root = r"E:\AllData\zunYiBreast"  # 数据的根路径，可以设置为“nii_path”的上一级目录

volume_path = r"E:\AllData\zunYiBreast\nii\volume"
seg_path = r"E:\AllData\zunYiBreast\nii\segmentation"

pretreatment_volume_path = os.path.join(data_root, 'cutCenter256/volume')
pretreatment_seg_path = os.path.join(data_root, 'cutCenter256/segmentation')


def identicSpacing(file,processfile_path,save_path):

    # 将CT和标签入读内存
    volume = sitk.ReadImage(os.path.join(processfile_path, file), sitk.sitkInt16)
    volume_array = sitk.GetArrayFromImage(volume)
    #96,
    target_size=256
    shape_x=448
    shape_y=448
    x_begin=shape_x//2-target_size//2
    x_end = shape_x // 2+target_size//2
    y_begin = shape_y // 2-target_size//2
    y_end = shape_y // 2+target_size//2

    volume_array = volume_array[:,x_begin:x_end,y_begin:y_end]
    '''
    后面要把spacing调好
    newspacing*newshape=oldshape*oldspacing
    '''
    sample_duration = 32
    sample_size = 256
    old_shape = volume_array.shape
    last_shape = np.array([sample_duration, sample_size, sample_size])  # 调整spacing之后的维度
    old_spacing = volume.GetSpacing()[2], volume.GetSpacing()[1], volume.GetSpacing()[0]
    old_spacing = np.array(old_spacing)
    dimension_adjustment = 1 / (old_shape / last_shape)
    volume_array_new = ndimage.zoom(volume_array, dimension_adjustment, order = 3)
    '''
    后面要把spacing调好
    newspacing*newshape=oldshape*oldspacing
    '''
    new_volume = sitk.GetImageFromArray(volume_array_new)
    new_spacing = old_shape * old_spacing / volume_array_new.shape
    new_volume.SetDirection(volume.GetDirection())
    new_volume.SetOrigin(volume.GetOrigin())
    new_volume.SetSpacing((new_spacing[2], new_spacing[1], new_spacing[0]))
    print(file, new_volume.GetSpacing(), volume_array_new.shape)
    sitk.WriteImage(new_volume, os.path.join(save_path, file))




if __name__ == '__main__':

    # if os.path.exists(para.pretreatment_path):
    #     shutil.rmtree(para.pretreatment_path)
    if not os.path.exists(pretreatment_volume_path):
        os.makedirs(pretreatment_volume_path)
    if not os.path.exists(pretreatment_seg_path):
        os.makedirs(pretreatment_seg_path)
    start = time()
    file_list=os.listdir(volume_path)
    print("处理MRI")
    for file in file_list:
        identicSpacing(file,processfile_path = volume_path,save_path = pretreatment_volume_path)

    file_list = os.listdir(seg_path)
    print("处理标注")
    for file in file_list:
        identicSpacing(file, processfile_path = seg_path, save_path = pretreatment_seg_path)



